#!/usr/bin/env python3
"""
DataPlexus 数据源管理性能验证脚本
验证API性能、数据库查询性能、并发处理能力
"""

import sys
import os
import asyncio
import time
import statistics
from datetime import datetime
from typing import List, Dict, Any

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, text
from app.core.database import get_async_session
from app.models.datasource import DataSource
from app.services.datasource_service import datasource_service
from app.schemas.datasource import DataSourceQueryParams
from app.core.logger import logger


class PerformanceVerifier:
    """性能验证器"""
    
    def __init__(self):
        self.test_results = []
    
    async def test_database_query_performance(self, db: AsyncSession) -> Dict[str, Any]:
        """测试数据库查询性能"""
        logger.info("测试数据库查询性能...")
        
        try:
            # 测试简单查询
            simple_query_times = []
            for i in range(10):
                start_time = time.time()
                
                stmt = select(DataSource).where(DataSource.is_deleted == False).limit(10)
                result = await db.execute(stmt)
                datasources = result.scalars().all()
                
                end_time = time.time()
                query_time = (end_time - start_time) * 1000  # 转换为毫秒
                simple_query_times.append(query_time)
                
                await asyncio.sleep(0.1)  # 避免过于频繁的查询
            
            # 测试复杂查询（带条件和排序）
            complex_query_times = []
            for i in range(5):
                start_time = time.time()
                
                stmt = select(DataSource).where(
                    DataSource.is_deleted == False
                ).order_by(DataSource.created_at.desc()).limit(20)
                result = await db.execute(stmt)
                datasources = result.scalars().all()
                
                end_time = time.time()
                query_time = (end_time - start_time) * 1000
                complex_query_times.append(query_time)
                
                await asyncio.sleep(0.1)
            
            # 计算统计信息
            simple_avg = statistics.mean(simple_query_times)
            simple_max = max(simple_query_times)
            complex_avg = statistics.mean(complex_query_times)
            complex_max = max(complex_query_times)
            
            # 性能标准：简单查询 < 50ms，复杂查询 < 100ms
            simple_pass = simple_avg < 50 and simple_max < 100
            complex_pass = complex_avg < 100 and complex_max < 200
            
            return {
                'success': simple_pass and complex_pass,
                'simple_query': {
                    'avg_time': simple_avg,
                    'max_time': simple_max,
                    'all_times': simple_query_times,
                    'pass': simple_pass
                },
                'complex_query': {
                    'avg_time': complex_avg,
                    'max_time': complex_max,
                    'all_times': complex_query_times,
                    'pass': complex_pass
                },
                'performance_standard': {
                    'simple_query_target': '< 50ms avg, < 100ms max',
                    'complex_query_target': '< 100ms avg, < 200ms max'
                }
            }
            
        except Exception as e:
            logger.error(f"数据库查询性能测试失败: {str(e)}")
            return {
                'success': False,
                'error': str(e)
            }
    
    async def test_api_response_performance(self, db: AsyncSession) -> Dict[str, Any]:
        """测试API响应性能"""
        logger.info("测试API响应性能...")
        
        try:
            # 模拟用户类（简化）
            class MockUser:
                def __init__(self):
                    self.id = 1
                    self.username = 'test_user'
            
            mock_user = MockUser()
            
            # 测试获取数据源列表的性能
            list_api_times = []
            for i in range(10):
                start_time = time.time()
                
                params = DataSourceQueryParams(page=1, page_size=10)
                result = await datasource_service.get_datasources(
                    db=db,
                    page=params.page,
                    page_size=params.page_size,
                    current_user=mock_user
                )
                
                end_time = time.time()
                api_time = (end_time - start_time) * 1000
                list_api_times.append(api_time)
                
                await asyncio.sleep(0.1)
            
            # 测试带搜索的API性能
            search_api_times = []
            for i in range(5):
                start_time = time.time()
                
                result = await datasource_service.get_datasources(
                    db=db,
                    page=1,
                    page_size=10,
                    search='test',
                    current_user=mock_user
                )
                
                end_time = time.time()
                api_time = (end_time - start_time) * 1000
                search_api_times.append(api_time)
                
                await asyncio.sleep(0.1)
            
            # 计算统计信息
            list_avg = statistics.mean(list_api_times)
            list_max = max(list_api_times)
            search_avg = statistics.mean(search_api_times)
            search_max = max(search_api_times)
            
            # 性能标准：列表API < 200ms，搜索API < 300ms
            list_pass = list_avg < 200 and list_max < 500
            search_pass = search_avg < 300 and search_max < 600
            
            return {
                'success': list_pass and search_pass,
                'list_api': {
                    'avg_time': list_avg,
                    'max_time': list_max,
                    'all_times': list_api_times,
                    'pass': list_pass
                },
                'search_api': {
                    'avg_time': search_avg,
                    'max_time': search_max,
                    'all_times': search_api_times,
                    'pass': search_pass
                },
                'performance_standard': {
                    'list_api_target': '< 200ms avg, < 500ms max',
                    'search_api_target': '< 300ms avg, < 600ms max'
                }
            }
            
        except Exception as e:
            logger.error(f"API响应性能测试失败: {str(e)}")
            return {
                'success': False,
                'error': str(e)
            }
    
    async def test_concurrent_request_performance(self, db: AsyncSession) -> Dict[str, Any]:
        """测试并发请求性能"""
        logger.info("测试并发请求性能...")
        
        try:
            class MockUser:
                def __init__(self):
                    self.id = 1
                    self.username = 'test_user'
            
            mock_user = MockUser()
            
            # 测试并发查询
            concurrent_levels = [5, 10, 20]
            results = {}
            
            for level in concurrent_levels:
                logger.info(f"测试 {level} 个并发请求...")
                
                async def single_request():
                    start_time = time.time()
                    result = await datasource_service.get_datasources(
                        db=db,
                        page=1,
                        page_size=5,
                        current_user=mock_user
                    )
                    end_time = time.time()
                    return (end_time - start_time) * 1000
                
                # 创建并发任务
                tasks = [single_request() for _ in range(level)]
                
                start_time = time.time()
                response_times = await asyncio.gather(*tasks, return_exceptions=True)
                total_time = (time.time() - start_time) * 1000
                
                # 过滤异常结果
                valid_times = [t for t in response_times if isinstance(t, (int, float))]
                error_count = len(response_times) - len(valid_times)
                
                if valid_times:
                    avg_response_time = statistics.mean(valid_times)
                    max_response_time = max(valid_times)
                    min_response_time = min(valid_times)
                else:
                    avg_response_time = max_response_time = min_response_time = 0
                
                # 性能标准：并发请求平均响应时间应该不超过单个请求的2倍
                success_rate = len(valid_times) / level
                performance_acceptable = avg_response_time < 1000 and success_rate > 0.8
                
                results[f'concurrent_{level}'] = {
                    'total_time': total_time,
                    'avg_response_time': avg_response_time,
                    'max_response_time': max_response_time,
                    'min_response_time': min_response_time,
                    'success_count': len(valid_times),
                    'error_count': error_count,
                    'success_rate': success_rate,
                    'performance_acceptable': performance_acceptable
                }
                
                await asyncio.sleep(1)  # 测试间隔
            
            # 计算总体性能
            overall_success = all(
                result['performance_acceptable'] 
                for result in results.values()
            )
            
            return {
                'success': overall_success,
                'concurrent_tests': results,
                'performance_standard': {
                    'target': '平均响应时间 < 1000ms, 成功率 > 80%'
                }
            }
            
        except Exception as e:
            logger.error(f"并发请求性能测试失败: {str(e)}")
            return {
                'success': False,
                'error': str(e)
            }
    
    async def test_memory_usage(self) -> Dict[str, Any]:
        """测试内存使用情况"""
        logger.info("测试内存使用情况...")
        
        try:
            import psutil
            import gc
            
            # 获取当前进程
            process = psutil.Process()
            
            # 记录初始内存使用
            initial_memory = process.memory_info().rss / 1024 / 1024  # MB
            
            # 模拟大量数据处理
            test_data = []
            for i in range(1000):
                test_data.append({
                    'id': i,
                    'name': f'test_datasource_{i}',
                    'data': 'x' * 1000  # 1KB数据
                })
            
            # 记录峰值内存使用
            peak_memory = process.memory_info().rss / 1024 / 1024  # MB
            
            # 清理数据
            del test_data
            gc.collect()
            
            # 记录清理后内存使用
            final_memory = process.memory_info().rss / 1024 / 1024  # MB
            
            memory_increase = peak_memory - initial_memory
            memory_recovered = peak_memory - final_memory
            recovery_rate = memory_recovered / memory_increase if memory_increase > 0 else 1
            
            # 性能标准：内存增长 < 100MB，回收率 > 80%
            memory_acceptable = memory_increase < 100 and recovery_rate > 0.8
            
            return {
                'success': memory_acceptable,
                'initial_memory_mb': initial_memory,
                'peak_memory_mb': peak_memory,
                'final_memory_mb': final_memory,
                'memory_increase_mb': memory_increase,
                'memory_recovered_mb': memory_recovered,
                'recovery_rate': recovery_rate,
                'performance_standard': {
                    'target': '内存增长 < 100MB, 回收率 > 80%'
                }
            }
            
        except ImportError:
            logger.warning("psutil未安装，跳过内存测试")
            return {
                'success': True,
                'skipped': True,
                'reason': 'psutil not installed'
            }
        except Exception as e:
            logger.error(f"内存使用测试失败: {str(e)}")
            return {
                'success': False,
                'error': str(e)
            }
    
    async def run_all_performance_tests(self):
        """运行所有性能测试"""
        logger.info("开始运行性能验证测试...")
        
        start_time = datetime.now()
        test_results = {
            'start_time': start_time,
            'database_query': None,
            'api_response': None,
            'concurrent_request': None,
            'memory_usage': None,
            'overall_success': False
        }
        
        try:
            async with get_async_session() as db:
                # 数据库查询性能测试
                test_results['database_query'] = await self.test_database_query_performance(db)
                
                # API响应性能测试
                test_results['api_response'] = await self.test_api_response_performance(db)
                
                # 并发请求性能测试
                test_results['concurrent_request'] = await self.test_concurrent_request_performance(db)
            
            # 内存使用测试
            test_results['memory_usage'] = await self.test_memory_usage()
            
            # 计算总体成功状态
            test_results['overall_success'] = all([
                test_results['database_query']['success'],
                test_results['api_response']['success'],
                test_results['concurrent_request']['success'],
                test_results['memory_usage']['success']
            ])
            
            test_results['end_time'] = datetime.now()
            test_results['duration'] = (test_results['end_time'] - test_results['start_time']).total_seconds()
            
            return test_results
            
        except Exception as e:
            logger.error(f"性能验证测试失败: {str(e)}")
            raise


def print_performance_results(results: Dict[str, Any]):
    """打印性能测试结果"""
    print("\n" + "="*80)
    print("⚡ DataPlexus 数据源管理性能验证报告")
    print("="*80)
    
    # 数据库查询性能
    db_query = results['database_query']
    print(f"\n🗄️ 数据库查询性能:")
    print(f"   结果: {'✅ 通过' if db_query['success'] else '❌ 失败'}")
    if db_query['success']:
        print(f"   简单查询平均: {db_query['simple_query']['avg_time']:.2f}ms")
        print(f"   复杂查询平均: {db_query['complex_query']['avg_time']:.2f}ms")
    else:
        print(f"   错误: {db_query.get('error', '未知错误')}")
    
    # API响应性能
    api_resp = results['api_response']
    print(f"\n🌐 API响应性能:")
    print(f"   结果: {'✅ 通过' if api_resp['success'] else '❌ 失败'}")
    if api_resp['success']:
        print(f"   列表API平均: {api_resp['list_api']['avg_time']:.2f}ms")
        print(f"   搜索API平均: {api_resp['search_api']['avg_time']:.2f}ms")
    else:
        print(f"   错误: {api_resp.get('error', '未知错误')}")
    
    # 并发请求性能
    concurrent = results['concurrent_request']
    print(f"\n🔄 并发请求性能:")
    print(f"   结果: {'✅ 通过' if concurrent['success'] else '❌ 失败'}")
    if concurrent['success']:
        for test_name, test_result in concurrent['concurrent_tests'].items():
            level = test_name.split('_')[1]
            print(f"   {level}并发: 平均{test_result['avg_response_time']:.2f}ms, 成功率{test_result['success_rate']*100:.1f}%")
    else:
        print(f"   错误: {concurrent.get('error', '未知错误')}")
    
    # 内存使用
    memory = results['memory_usage']
    print(f"\n💾 内存使用:")
    print(f"   结果: {'✅ 通过' if memory['success'] else '❌ 失败'}")
    if memory['success'] and not memory.get('skipped'):
        print(f"   内存增长: {memory['memory_increase_mb']:.2f}MB")
        print(f"   回收率: {memory['recovery_rate']*100:.1f}%")
    elif memory.get('skipped'):
        print(f"   跳过: {memory['reason']}")
    else:
        print(f"   错误: {memory.get('error', '未知错误')}")
    
    # 总体结果
    print(f"\n🎯 总体性能评估:")
    print(f"   测试耗时: {results['duration']:.2f}秒")
    print(f"   总体状态: {'✅ 性能优秀' if results['overall_success'] else '⚠️ 需要优化'}")
    
    print("\n" + "="*80)


async def main():
    """主函数"""
    try:
        verifier = PerformanceVerifier()
        results = await verifier.run_all_performance_tests()
        
        print_performance_results(results)
        
        return results['overall_success']
        
    except Exception as e:
        print(f"\n❌ 性能验证失败: {str(e)}")
        return False


if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)
